# Copyright 2024 Marimo. All rights reserved.
from __future__ import annotations

import threading
from collections import deque
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, cast

from marimo import _loggers
from marimo._ast.cell import CellImpl
from marimo._ast.compiler import code_key
from marimo._ast.sql_visitor import SQLRef, SQLTypes
from marimo._ast.variables import is_mangled_local
from marimo._ast.visitor import ImportData, Name, VariableData
from marimo._runtime.executor import (
    ExecutionConfig,
    get_executor,
)
from marimo._types.ids import CellId_t

if TYPE_CHECKING:
    from collections.abc import Collection


Edge = tuple[CellId_t, CellId_t]
# EdgeWithVar uses a list rather than a set for the variables linking the cells
# as sets are not JSON-serializable (required by static_notebook_template()).
# The first entry is the source node; the second entry is a list of defs from
# the source read by the destination; and the third entry is the destination
# node.
EdgeWithVar = tuple[CellId_t, tuple[str, ...], CellId_t]

LOGGER = _loggers.marimo_logger()


# TODO(akshayka): Add method disable_cell, enable_cell which handle
# state transitions on cells
@dataclass(frozen=True)
class DirectedGraph:
    # Nodes in the graph
    cells: dict[CellId_t, CellImpl] = field(default_factory=dict)

    # Edge (u, v) means v is a child of u, i.e., v has a reference
    # to something defined in u
    children: dict[CellId_t, set[CellId_t]] = field(default_factory=dict)

    # Reversed edges (parent pointers) for convenience
    parents: dict[CellId_t, set[CellId_t]] = field(default_factory=dict)

    # Cells that define the same name
    #
    # siblings[cell_id] is a set of cell ids, one for each cell that shares a
    # definition with cell_id.
    #
    # If this dict is non-empty, then the marimo program contains multiply
    # defined names (and is therefore in an error state)
    siblings: dict[CellId_t, set[CellId_t]] = field(default_factory=dict)

    # A mapping from defs to the cells that define them
    definitions: dict[Name, set[CellId_t]] = field(default_factory=dict)
    typed_definitions: dict[tuple[Name, str], set[CellId_t]] = field(
        default_factory=dict
    )
    definition_types: dict[Name, set[str]] = field(default_factory=dict)

    # The set of cycles in the graph
    cycles: set[tuple[Edge, ...]] = field(default_factory=set)

    # This lock must be acquired during methods that mutate the graph; it's
    # only needed because a graph is shared between the kernel and the code
    # completion service. It should almost always be uncontended.
    lock: threading.Lock = field(default_factory=threading.Lock)

    def is_cell_cached(self, cell_id: CellId_t, code: str) -> bool:
        """Whether a cell with id `cell_id` and code `code` is in the graph."""
        return (
            cell_id in self.cells and code_key(code) == self.cells[cell_id].key
        )

    # TODO: language type?
    def get_defining_cells(self, name: Name) -> set[CellId_t]:
        """Get all cells that define name.

        This is a singleton for well-formed graphs.
        """
        return self.definitions[name]

    def get_referring_cells(
        self, name: Name, language: Literal["python", "sql"]
    ) -> set[CellId_t]:
        """Get all cells that have a ref to `name`.

        The variable can be either a Python variable or a SQL variable (table).
        SQL variables don't leak to Python cells, but Python variables do leak
        to SQL.

        Only does a local analysis of refs, without taking into consideration
        whether refs are defined by other cells.
        """
        if language == "sql":
            # For SQL, only return SQL cells that reference the name
            cells = set()
            for cid, cell in self.cells.items():
                if cell.language != "sql":
                    continue

                for ref in cell.refs:
                    # Direct reference match
                    if ref == name:
                        cells.add(cid)
                        break

                    sql_ref = cell.sql_refs.get(ref)

                    kind: SQLTypes = "any"
                    if name in cell.variable_data:
                        variable_data = cell.variable_data[name][-1]
                        kind = cast(SQLTypes, variable_data.kind)

                    # Hierarchical reference match
                    if sql_ref and sql_ref.matches_hierarchical_ref(
                        name, ref, kind
                    ):
                        cells.add(cid)
                        break

            return cells
        else:
            # For Python, return all cells that reference the name
            return {
                cid for cid, cell in self.cells.items() if name in cell.refs
            }

    def _find_sql_hierarchical_matches(
        self, sql_ref: SQLRef
    ) -> list[tuple[set[CellId_t], Name]]:
        """
        This method searches through all definitions in the graph to find cells
        that define the individual components (table, schema, or catalog) of the
        hierarchical reference.

        For example, given a reference "my_schema.my_table", this method will:
        - Look for cells that define a table/view named "my_table"
        - Look for cells that define a catalog named "my_schema"
          (when the reference has at least 2 parts)

        Args:
            sql_ref: A hierarchical SQL reference (e.g., "schema.table",
                  "catalog.schema.table") to find matching definitions for.

        Returns:
            A list of tuples containing:
            - A set of cell IDs that define components of the hierarchical reference
            - The definition of the name that was found (e.g., "schema.table" -> "table")
        """
        matching_cell_ids_list = []

        for (def_name, kind), cell_ids in self.typed_definitions.items():
            # Match table/view definitions
            if sql_ref.contains_hierarchical_ref(def_name, kind):
                matching_cell_ids_list.append((cell_ids, def_name))

        return matching_cell_ids_list

    def _is_valid_cell_reference(
        self, cell_id: CellId_t, variable_name: Name
    ) -> bool:
        """Check if a cell reference is valid and log errors if not."""
        if cell_id not in self.cells:
            LOGGER.error(
                "Variable %s is defined in cell %s, but is not in the graph",
                variable_name,
                cell_id,
            )
            return False
        return True

    def _resolve_variable_name(
        self,
        name: Name,
        other_cell: CellImpl,
        sql_ref: Optional[SQLRef],
        sql_matches: list[tuple[set[CellId_t], Name]],
    ) -> Name:
        """
        Resolve the variable name to use when checking if it exists in another cell.

        For regular (non-SQL) references, returns the original name unchanged.
        For SQL hierarchical references, finds the variable name from sql_matches that
        is actually defined in the other_cell.

        Example:
            cell_1: CREATE SCHEMA schema_name
            cell_2: CREATE TABLE schema_name.table_name
            cell_3: FROM schema_name.table_name SELECT *

            When cell_3 references "schema_name.table_name":
            - For cell_1: returns "schema_name" (what cell_1 actually defines)
            - For cell_2: returns "table_name" (what cell_2 actually defines)
        """
        if not sql_ref or name in other_cell.variable_data:
            return name

        # For SQL hierarchical references, find the matching variable name
        for _, matching_variable_name in sql_matches:
            if matching_variable_name in other_cell.variable_data:
                return matching_variable_name

        return name

    def get_path(self, source: CellId_t, dst: CellId_t) -> list[Edge]:
        """Get a path from `source` to `dst`, if any."""
        if source == dst:
            return []

        # deque has O(1) append/pop operation
        queue: deque[tuple[CellId_t, list[Edge]]] = deque([(source, [])])
        found = {source}  # set has O(1) lookups

        while queue:
            node, path = queue.popleft()  # O(1) operation
            for cid in self.children[node]:
                if cid not in found:
                    next_path = path + [(node, cid)]
                    if cid == dst:
                        return next_path
                    found.add(cid)
                    queue.append((cid, next_path))
        return []

    def register_cell(self, cell_id: CellId_t, cell: CellImpl) -> None:
        """Add a cell to the graph.

        Mutates the graph, acquiring `self.lock`.

        Requires that `cell_id` is not already in the graph.
        """
        LOGGER.debug("Acquiring graph lock to register cell %s", cell_id)
        with self.lock:
            LOGGER.debug("Acquired graph lock.")
            assert cell_id not in self.cells
            self.cells[cell_id] = cell
            # Children are the set of cells that refer to a name defined in
            # `cell`
            children: set[CellId_t] = set()
            # Cells that define the same name as this one
            siblings: set[CellId_t] = set()
            # Parents are the set of cells that define a name referred to by
            # `cell`
            parents: set[CellId_t] = set()

            # Populate children, siblings, and parents
            self.children[cell_id] = children
            self.siblings[cell_id] = siblings
            self.parents[cell_id] = parents

            # First, we process the variables that this cell defines. Any cell
            # that refers to a defined variable becomes a child of this cell;
            # any cell that defines a variable defined by this cell becomes
            # a sibling.
            for name, variable_data in cell.variable_data.items():
                # NB. Only the last definition matters.
                # Technically more nuanced with branching statements, but this is
                # the best we can do with static analysis.
                variable = variable_data[-1]
                typed_def = (name, variable.kind)

                if (
                    name in self.definitions
                    and typed_def not in self.typed_definitions
                ):
                    # Duplicate if the qualified name is no different
                    if (
                        variable.qualified_name == name
                        or variable.language != "sql"
                    ):
                        self.definitions[name].add(cell_id)
                else:
                    self.definitions.setdefault(name, set()).add(cell_id)
                self.typed_definitions.setdefault(typed_def, set()).add(
                    cell_id
                )
                self.definition_types.setdefault(name, set()).add(
                    variable.kind
                )

                for sibling in self.definitions[name]:
                    # TODO(akshayka): Distinguish between Python/SQL?
                    if sibling != cell_id:
                        siblings.add(sibling)
                        self.siblings[sibling].add(cell_id)

                # a cell can refer to its own defs, but that doesn't add an
                # edge to the dependency graph
                referring_cells = self.get_referring_cells(
                    name,
                    language=variable_data[-1].language,
                ) - set((cell_id,))
                # we will add an edge (cell_id, v) for each v in
                # referring_cells; if there is a path from v to cell_id, then
                # the new edge will form a cycle
                for v in referring_cells:
                    path = self.get_path(v, cell_id)
                    if path:
                        self.cycles.add(tuple([(cell_id, v)] + path))

                children.update(referring_cells)
                for child in referring_cells:
                    self.parents[child].add(cell_id)

            # Next, we process the cells references. The cell becomes a child
            # of cells that define its referenced variables. We also have
            # special logic for handling references that are deleted by this cell,
            # since cells that delete variables that were defined elsewhere
            # are made children of cells that reference that variable.

            for name in cell.refs:
                # First, for each referenced variable, we add cells that define
                # that variable as parents
                other_ids_defining_name: set[CellId_t] = (
                    self.definitions[name]
                    if name in self.definitions
                    else set()
                ) - set((cell_id,))

                variable_name: Name = name

                # Handle SQL matching for hierarchical references
                sql_matches: list[tuple[set[CellId_t], Name]] = []
                sql_ref = cell.sql_refs.get(name)
                if sql_ref:
                    sql_matches = self._find_sql_hierarchical_matches(sql_ref)
                    for matching_cell_ids, _ in sql_matches:
                        if cell_id in matching_cell_ids:
                            LOGGER.debug(
                                "Cell %s is referencing itself", cell_id
                            )
                            continue
                        other_ids_defining_name.update(matching_cell_ids)

                # If other_ids_defining_name is empty, the user will get a
                # NameError at runtime (unless the symbol is a builtin).
                for other_id in other_ids_defining_name:
                    if other_id == cell_id:
                        LOGGER.error("Cell %s is referencing itself", cell_id)
                        continue
                    if not self._is_valid_cell_reference(
                        other_id, variable_name
                    ):
                        continue
                    other_cell = self.cells[other_id]

                    variable_name = self._resolve_variable_name(
                        variable_name, other_cell, sql_ref, sql_matches
                    )

                    # If we don't have a matching variable name, skip
                    if variable_name not in other_cell.variable_data:
                        LOGGER.error(
                            "Variable %s is not defined in cell %s",
                            variable_name,
                            other_id,
                        )
                        continue

                    other_variable_data = other_cell.variable_data[
                        variable_name
                    ][-1]
                    language = other_variable_data.language
                    if language == "sql" and cell.language == "python":
                        # SQL table/db def -> Python ref is not an edge
                        continue
                    if language == "sql" and cell.language == "sql":
                        # Edges between SQL cells need to respect hierarchy.
                        if sql_ref and not sql_ref.matches_hierarchical_ref(
                            variable_name,
                            other_variable_data.qualified_name or name,
                            kind=cast(SQLTypes, other_variable_data.kind),
                        ):
                            continue
                    parents.add(other_id)
                    # we are adding an edge (other_id, cell_id). If there
                    # is a path from cell_id to other_id, then the new
                    # edge forms a cycle
                    path = self.get_path(cell_id, other_id)
                    if path:
                        self.cycles.add(tuple([(other_id, cell_id)] + path))
                    self.children[other_id].add(cell_id)

                # Next, any cell that deletes this referenced variable is made
                # a child of this cell. In particular, if a cell deletes a
                # variable, it becomes a child of all other cells that
                # reference that variable. This means that if two cells delete
                # the same variable, they form a cycle.
                #
                # For example, two cells
                #
                #   cell u: x
                #   cell v: del x
                #
                # v becomes a child of u.
                #
                # Another example:
                #
                #   cell u: del x
                #   cell v: del x
                #
                # u and v form a cycle.
                other_ids_deleting_name: set[CellId_t] = set(
                    cid
                    for cid in self.get_referring_cells(
                        name, language="python"
                    )
                    if name in self.cells[cid].deleted_refs
                ) - set((cell_id,))
                for v in other_ids_deleting_name:
                    path = self.get_path(v, cell_id)
                    if path:
                        self.cycles.add(tuple([(cell_id, v)] + path))
                    self.parents[v].add(cell_id)
                children.update(other_ids_deleting_name)

            # Finally, if this cell deletes a variable, we make it a child of
            # all other cells that reference this variable.
            for name in cell.deleted_refs:
                referring_cells = self.get_referring_cells(
                    name, language="python"
                ) - set((cell_id,))
                for other_id in referring_cells:
                    parents.add(other_id)
                    path = self.get_path(cell_id, other_id)
                    if path:
                        self.cycles.add(tuple([(other_id, cell_id)] + path))
                    self.children[other_id].add(cell_id)

        LOGGER.debug("Registered cell %s and released graph lock", cell_id)
        if self.is_any_ancestor_stale(cell_id):
            self.set_stale(set([cell_id]))

        if self.is_any_ancestor_disabled(cell_id):
            cell.set_runtime_state(status="disabled-transitively")

    def is_any_ancestor_stale(self, cell_id: CellId_t) -> bool:
        return any(self.cells[cid].stale for cid in self.ancestors(cell_id))

    def is_any_ancestor_disabled(self, cell_id: CellId_t) -> bool:
        return any(
            self.cells[cid].config.disabled for cid in self.ancestors(cell_id)
        )

    def disable_cell(self, cell_id: CellId_t) -> None:
        """
        Disables a cell in the graph.

        Does not mutate the graph (but does mutate cell statuses).

        Returns the ids of descendants that are disabled transitively.
        """
        if cell_id not in self.cells:
            raise ValueError(f"Cell {cell_id} not found")

        for cid in transitive_closure(self, set([cell_id])) - set([cell_id]):
            cell = self.cells[cid]
            cell.set_runtime_state(status="disabled-transitively")

    def enable_cell(self, cell_id: CellId_t) -> set[CellId_t]:
        """
        Enables a cell in the graph.

        Does not mutate the graph (but does mutate cell statuses).

        Returns:
        - set of cells that were stale and should be re-run
        """
        if cell_id not in self.cells:
            raise ValueError(f"Cell {cell_id} not found")

        cells_to_run: set[CellId_t] = set()
        for cid in transitive_closure(self, set([cell_id])):
            if not self.is_disabled(cid):
                child = self.cells[cid]
                if child.stale:
                    # cell was previously disabled, is no longer
                    # disabled, and is stale: needs to run.
                    cells_to_run.add(cid)
                if child.disabled_transitively:
                    # cell is no longer disabled: status -> idle
                    child.set_runtime_state("idle")
        return cells_to_run

    def delete_cell(self, cell_id: CellId_t) -> set[CellId_t]:
        """Removes a cell from the graph.

        Mutates the graph, acquiring `self.lock`.

        Returns the ids of the children of the removed cell.
        """
        LOGGER.debug("Acquiring graph lock to delete cell %s", cell_id)
        with self.lock:
            LOGGER.debug("Acquired graph lock to delete cell %s", cell_id)
            if cell_id not in self.cells:
                raise ValueError(f"Cell {cell_id} not found")

            # Removing this cell from its defs' definer sets
            for name in self.cells[cell_id].defs:
                name_defs = self.definitions[name]
                name_defs.remove(cell_id)
                if not name_defs:
                    # No more cells define this name, so we remove it from the
                    # graph
                    del self.definitions[name]
                    # Clean up all typed definitions
                    for typed_def in self.definition_types[name]:
                        del self.typed_definitions[(name, typed_def)]
                    del self.definition_types[name]

            # Remove cycles that are broken from removing this cell.
            edges = [(cell_id, child) for child in self.children[cell_id]] + [
                (parent, cell_id) for parent in self.parents[cell_id]
            ]
            for e in edges:
                broken_cycles = [c for c in self.cycles if e in c]
                for c in broken_cycles:
                    self.cycles.remove(c)

            # Grab a reference to children before we remove it from our map.
            children = self.children[cell_id]

            # Purge this cell from the graph.
            del self.cells[cell_id]
            del self.children[cell_id]
            del self.parents[cell_id]
            del self.siblings[cell_id]

            for elems in self.parents.values():
                if cell_id in elems:
                    elems.remove(cell_id)
            for elems in self.children.values():
                if cell_id in elems:
                    elems.remove(cell_id)
            for elems in self.siblings.values():
                if cell_id in elems:
                    elems.remove(cell_id)
        LOGGER.debug("Deleted cell %s and Released graph lock.", cell_id)
        return children

    def is_disabled(self, cell_id: CellId_t) -> bool:
        if cell_id not in self.cells:
            raise ValueError(f"Cell {cell_id} not in graph.")
        cell = self.cells[cell_id]
        if cell.config.disabled:
            return True
        seen: set[CellId_t] = set()
        queue = [cell_id]
        while queue:
            cid = queue.pop()
            seen.add(cid)
            for parent_id in self.parents[cid]:
                if parent_id in seen:
                    continue
                elif self.cells[parent_id].config.disabled:
                    return True
                else:
                    queue.append(parent_id)
        return False

    def get_imports(
        self, cell_id: Optional[CellId_t] = None
    ) -> dict[Name, ImportData]:
        imports = {}
        cells = (
            self.cells.values() if cell_id is None else [self.cells[cell_id]]
        )
        for cell in cells:
            for imported in cell.imports:
                imports[imported.definition] = imported
        return imports

    def get_multiply_defined(self) -> list[Name]:
        """Return a list of names that are defined in multiple cells"""
        names: list[Name] = []
        for name, definers in self.definitions.items():
            if len(definers) > 1:
                names.append(name)
        return names

    def get_deleted_nonlocal_ref(self) -> list[Name]:
        names: list[Name] = []
        for cell in self.cells.values():
            for ref in cell.deleted_refs:
                if ref in self.definitions:
                    names.append(ref)
        return names

    def descendants(self, cell_id: CellId_t) -> set[CellId_t]:
        return transitive_closure(self, set([cell_id]), inclusive=False)

    def ancestors(self, cell_id: CellId_t) -> set[CellId_t]:
        return transitive_closure(
            self, set([cell_id]), children=False, inclusive=False
        )

    def set_stale(
        self, cell_ids: set[CellId_t], prune_imports: bool = False
    ) -> None:
        relatives = None if not prune_imports else import_block_relatives
        for cid in transitive_closure(self, cell_ids, relatives=relatives):
            self.cells[cid].set_stale(stale=True)

    def get_stale(self) -> set[CellId_t]:
        return set([cid for cid, cell in self.cells.items() if cell.stale])

    def get_transitive_references(
        self,
        refs: set[Name],
        inclusive: bool = True,
        predicate: Callable[[Name, VariableData], bool] | None = None,
    ) -> set[Name]:
        """Return a set of the passed-in cells' references and their
        references on the block (function / class) level.

        If inclusive, includes the references of the passed-in cells in the
        set.

        If predicate, only references satisfying predicate(ref) are included
        """
        # TODO: Consider caching on the graph level and updating on register /
        # delete
        processed: set[Name] = set()
        queue: set[Name] = refs & self.definitions.keys()
        predicate = predicate or (lambda *_: True)

        while queue:
            # Should ideally be one cell per ref, but for completion, stay
            # agnostic to potenital cycles.
            cells = {
                cell_id
                for ref in queue
                for cell_id in self.definitions.get(ref, set())
            }

            for cell_id in cells:
                data = self.cells[cell_id].variable_data
                newly_processed = set(data.keys()) & queue
                processed.update(newly_processed)
                queue.difference_update(newly_processed)
                for variable in newly_processed:
                    # variables can be defined multiple times in a single
                    # cell ...
                    for datum in data[variable]:
                        if predicate(variable, datum):
                            to_process = datum.required_refs - processed
                            queue.update(to_process & self.definitions.keys())
                            # Private variables referenced by public functions
                            # have to be included.
                            processed.update(
                                maybe_private
                                for maybe_private in (
                                    to_process - self.definitions.keys()
                                )
                                if is_mangled_local(maybe_private, cell_id)
                            )

        if inclusive:
            return processed | refs
        return processed - refs

    def copy(self, filename: None | str = None) -> DirectedGraph:
        """Return a deep copy of the graph by recompiling all cells.

        This is mainly useful in the case where recompilation must be done
        due to a dynamically changing notebook, where the line cache must be
        consistent with the cell code, e.g. for debugging.
        """
        from marimo._ast.compiler import compile_cell

        graph = DirectedGraph()
        with self.lock:
            for cid, old_cell in self.cells.items():
                cell = compile_cell(
                    old_cell.code,
                    cell_id=cid,
                    filename=filename,
                )
                # Carry over import data manually
                imported_defs = old_cell.import_workspace.imported_defs
                is_import_block = old_cell.import_workspace.is_import_block
                cell.import_workspace.imported_defs = imported_defs
                cell.import_workspace.is_import_block = is_import_block
                # Reregister
                graph.register_cell(cid, cell)
        return graph


def transitive_closure(
    graph: DirectedGraph,
    cell_ids: set[CellId_t],
    children: bool = True,
    inclusive: bool = True,
    relatives: (
        Callable[[DirectedGraph, CellId_t, bool], set[CellId_t]] | None
    ) = None,
    predicate: Callable[[CellImpl], bool] | None = None,
) -> set[CellId_t]:
    """Return a set of the passed-in cells and their descendants or ancestors

    If children is True, returns descendants; otherwise, returns ancestors

    If inclusive, includes passed-in cells in the set.

    If relatives is not None, it computes the parents/children of a
        cell

    If predicate, only cells satisfying predicate(cell) are included; applied
        after the relatives are computed
    """

    result: set[CellId_t] = cell_ids.copy() if inclusive else set()
    seen: set[CellId_t] = cell_ids.copy()
    queue: deque[CellId_t] = deque(cell_ids)
    predicate = predicate or (lambda _: True)

    def _relatives(cid: CellId_t) -> set[CellId_t]:
        if relatives is None:
            return graph.children[cid] if children else graph.parents[cid]
        return relatives(graph, cid, children)

    while queue:
        cid = queue.popleft()  # O(1) operation

        relatives_set = _relatives(cid)
        new_relatives = relatives_set - seen

        if new_relatives:
            # Add new relatives to queue and result if they pass predicate
            for relative in new_relatives:
                if predicate(graph.cells[relative]):
                    result.add(relative)
                seen.add(relative)
                queue.append(relative)

    return result


def induced_subgraph(
    graph: DirectedGraph, cell_ids: Collection[CellId_t]
) -> tuple[dict[CellId_t, set[CellId_t]], dict[CellId_t, set[CellId_t]]]:
    """Return parents and children for each node in `cell_ids`

    Represents the subgraph induced by `cell_ids`.
    """
    parents: dict[CellId_t, set[CellId_t]] = {}
    children: dict[CellId_t, set[CellId_t]] = {}
    for cid in cell_ids:
        parents[cid] = set(p for p in graph.parents[cid] if p in cell_ids)
        children[cid] = set(c for c in graph.children[cid] if c in cell_ids)
    return parents, children


def get_cycles(
    graph: DirectedGraph, cell_ids: Collection[CellId_t]
) -> list[tuple[Edge, ...]]:
    """Get all cycles among `cell_ids`."""
    _, induced_children = induced_subgraph(graph, cell_ids)
    induced_edges = set(
        [(u, v) for u in induced_children for v in induced_children[u]]
    )
    return [c for c in graph.cycles if all(e in induced_edges for e in c)]


def topological_sort(
    graph: DirectedGraph, cell_ids: Collection[CellId_t]
) -> list[CellId_t]:
    """Sort `cell_ids` in a topological order using a heap queue.

    When multiple cells have the same parents (including no parents), the tie is broken by
    registration order - cells registered earlier are processed first.
    """
    from heapq import heapify, heappop, heappush

    # Use a list for O(1) lookup of registration order
    registration_order = list(graph.cells.keys())
    top_down_keys = {key: idx for idx, key in enumerate(registration_order)}

    # Build adjacency lists and in-degree counts
    parents, children = induced_subgraph(graph, cell_ids)
    in_degree = {cid: len(parents[cid]) for cid in cell_ids}

    # Initialize heap with roots
    heap = [
        (top_down_keys[cid], cid) for cid in cell_ids if in_degree[cid] == 0
    ]
    heapify(heap)

    sorted_cell_ids: list[CellId_t] = []
    while heap:
        _, cid = heappop(heap)
        sorted_cell_ids.append(cid)

        # Process children
        for child in children[cid]:
            in_degree[child] -= 1
            if in_degree[child] == 0:
                heappush(heap, (top_down_keys[child], child))

    return sorted_cell_ids


def import_block_relatives(
    graph: DirectedGraph, cid: CellId_t, children: bool
) -> set[CellId_t]:
    if not children:
        return graph.parents[cid]

    cell = graph.cells[cid]
    if not cell.import_workspace.is_import_block:
        return graph.children[cid]

    # This cell is an import block, which should be special cased:
    #
    # We prune definitions that have already been imported from the set of
    # definitions used to find the descendants of this cell.
    unimported_defs = cell.defs - cell.import_workspace.imported_defs

    children_ids = {
        child_id
        for name in unimported_defs
        for child_id in graph.get_referring_cells(name, language="python")
    }

    # If children haven't been executed, then still use imported defs;
    # handle an edge case when an import cell is interrupted by an
    # exception or user interrupt, so that a module is imported but the
    # cell's children haven't run.
    if cell.import_workspace.imported_defs:
        interrupted_states = {"interrupted", "cancelled", "marimo-error", None}
        children_ids.update(
            child_id
            for name in cell.import_workspace.imported_defs
            for child_id in graph.get_referring_cells(name, language="python")
            if graph.cells[child_id].run_result_status in interrupted_states
        )

    return children_ids


class Runner:
    """Utility for running individual cells in a graph

    This class provides methods to a run a cell in the graph and obtain its
    output (last expression) and the values of its defs.

    If needed, the runner will recursively compute the values of the cell's
    refs by executing its ancestors. Refs can also be substituted by the
    caller.

    TODO(akshayka): Add an API for caching defs across cell runs.
    """

    def __init__(self, graph: DirectedGraph) -> None:
        self._graph = graph
        self._executor = get_executor(ExecutionConfig())

    @staticmethod
    def _returns(cell_impl: CellImpl, glbls: dict[str, Any]) -> dict[str, Any]:
        return {name: glbls[name] for name in cell_impl.defs if name in glbls}

    @staticmethod
    def _substitute_refs(
        cell_impl: CellImpl,
        glbls: dict[str, Any],
        kwargs: dict[str, Any],
    ) -> None:
        for argname, argvalue in kwargs.items():
            if argname in cell_impl.refs:
                glbls[argname] = argvalue
            else:
                raise ValueError(
                    f"Cell got unexpected argument {argname}"
                    f"The allowed arguments are {cell_impl.refs}."
                )

    def _get_ancestors(
        self, cell_impl: CellImpl, kwargs: dict[str, Any]
    ) -> set[CellId_t]:
        # Get the transitive closure of parents defining unsubstituted refs
        graph = self._graph
        substitutions = set(kwargs.keys())
        unsubstituted_refs = cell_impl.refs - substitutions
        parent_ids = set(
            [
                parent_id
                for parent_id in graph.parents[cell_impl.cell_id]
                if graph.cells[parent_id].defs.intersection(unsubstituted_refs)
            ]
        )
        return transitive_closure(graph, parent_ids, children=False)

    @staticmethod
    def _validate_kwargs(cell_impl: CellImpl, kwargs: dict[str, Any]) -> None:
        for argname in kwargs:
            if argname not in cell_impl.refs:
                raise ValueError(
                    f"Cell got unexpected argument {argname}; "
                    f"The allowed arguments are {cell_impl.refs}."
                )

    def is_coroutine(self, cell_id: CellId_t) -> bool:
        return self._graph.cells[cell_id].is_coroutine() or any(
            self._graph.cells[cid].is_coroutine()
            for cid in self._get_ancestors(
                self._graph.cells[cell_id], kwargs={}
            )
        )

    async def run_cell_async(
        self, cell_id: CellId_t, kwargs: dict[str, Any]
    ) -> tuple[Any, dict[str, Any]]:
        """Run a possibly async cell and its ancestors

        Substitutes kwargs as refs for the cell, omitting ancestors that
        whose refs are substituted.
        """
        graph = self._graph
        cell_impl = graph.cells[cell_id]
        Runner._validate_kwargs(cell_impl, kwargs)
        ancestor_ids = self._get_ancestors(cell_impl, kwargs)

        glbls: dict[str, Any] = {}
        for cid in topological_sort(graph, ancestor_ids):
            await self._executor.execute_cell_async(
                graph.cells[cid], glbls, graph
            )

        Runner._substitute_refs(cell_impl, glbls, kwargs)
        output = await self._executor.execute_cell_async(
            graph.cells[cell_impl.cell_id], glbls, graph
        )
        defs = Runner._returns(cell_impl, glbls)
        return output, defs

    def run_cell_sync(
        self, cell_id: CellId_t, kwargs: dict[str, Any]
    ) -> tuple[Any, dict[str, Any]]:
        """Run a synchronous cell and its ancestors

        Substitutes kwargs as refs for the cell, omitting ancestors that
        whose refs are substituted.

        Raises a `RuntimeError` if the cell or any of its unsubstituted
        ancestors are coroutine functions.
        """
        graph = self._graph
        cell_impl = graph.cells[cell_id]
        if cell_impl.is_coroutine():
            raise RuntimeError(
                "A coroutine function can't be run synchronously. "
                "Use `run_async()` instead"
            )

        Runner._validate_kwargs(cell_impl, kwargs)
        ancestor_ids = self._get_ancestors(cell_impl, kwargs)

        if any(graph.cells[cid].is_coroutine() for cid in ancestor_ids):
            raise RuntimeError(
                "Cell has an ancestor that is a "
                "coroutine (async) cell. Use `run_async()` instead"
            )

        glbls: dict[str, Any] = {}
        for cid in topological_sort(graph, ancestor_ids):
            self._executor.execute_cell(graph.cells[cid], glbls, graph)

        self._substitute_refs(cell_impl, glbls, kwargs)
        output = self._executor.execute_cell(
            graph.cells[cell_impl.cell_id], glbls, graph
        )
        defs = Runner._returns(cell_impl, glbls)
        return output, defs
